package org.ws.task.netty.handler;

import com.google.common.collect.Maps;
import io.netty.channel.ChannelHandlerContext;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.curator.shaded.com.google.common.collect.Lists;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.ws.task.exception.RequestHandlerException;
import org.ws.task.item.TaskInstance;
import org.ws.task.repository.NodeMetadataRepository;
import org.ws.task.repository.TaskRepository;
import org.ws.task.schedule.DefaultItemSchedule;
import org.ws.task.schedule.ItemSchedule;
import org.ws.task.schedule.ScheduleStatus;
import org.ws.task.model.protocol.PullItemRequest;
import org.ws.task.model.protocol.PullItemResponse;
import org.ws.task.model.TaskItem;
import org.ws.task.netty.RequestActionHandler;
import org.ws.task.model.protocol.base.RequestAction;
import org.ws.task.model.protocol.base.RequestMessage;
import org.ws.task.model.protocol.base.ResponseMessage;
import org.ws.task.model.protocol.base.ResponseStatus;

import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

public class PullItemProviderHandler implements RequestActionHandler<PullItemRequest, PullItemResponse> {

    private static Logger log = LoggerFactory.getLogger(PullItemProviderHandler.class.getName());

    private static final int MAX_SIZE = 100;
    private final Map<String, AtomicInteger> indexMap;
    private final ItemSchedule<TaskItem> itemSchedule;

    public PullItemProviderHandler() {
        this.indexMap = Maps.newConcurrentMap();
        this.itemSchedule = new DefaultItemSchedule();
    }

    @Override
    public RequestAction action() {
        return RequestAction.PULL_TASK;
    }

    @Override
    public ResponseMessage<PullItemResponse> handler(ChannelHandlerContext ctx, RequestMessage<PullItemRequest> requestMessage) throws RequestHandlerException {
        int id = requestMessage.getId();
        ResponseMessage<PullItemResponse> responseMessage = new ResponseMessage<>(id);
        PullItemRequest pullItemRequest = requestMessage.getBodyObject();
        String taskName = pullItemRequest.getTaskName();
        String nodeId = pullItemRequest.getNodeId();
        Integer size = pullItemRequest.getSize();
        TaskInstance taskInstance = TaskRepository.getInstance().getTaskInstance(taskName);
        if(taskInstance != null){
            if(!indexMap.containsKey(taskName)){
                indexMap.put(taskName,new AtomicInteger(0));
            }
            AtomicInteger index = indexMap.get(taskName);
            if(size <= 0){
                size = taskInstance.getTaskConfig().getChunk();
            }
            if(size > MAX_SIZE){
                size = MAX_SIZE;
            }
            // 验证匹配
            if(NodeMetadataRepository.getInstance().isMatch(taskName,nodeId)){
                List<TaskItem> items = Lists.newArrayList();
                this.itemSchedule.ready(taskName,nodeId);
                final int chunk = size;
                final int scheduleSize = size * NodeMetadataRepository.getInstance().getWorkers().size();
                final AtomicInteger count = new AtomicInteger(0);
                while (this.itemSchedule.isContinue(taskName,nodeId)){
                    this.itemSchedule.schedule(taskName,nodeId,scheduleSize,item -> {
                        if(item != null){
                            items.add(item);
                            if(count.incrementAndGet()>=chunk){
                                return ScheduleStatus.STOP;
                            }else {
                                return ScheduleStatus.CONTINUABLE;
                            }
                        }
                        else {
                            return ScheduleStatus.FINISHED;
                        }
                    });
                }
                try {
                    PullItemResponse pullItemResponse = new PullItemResponse();
                    if(CollectionUtils.isNotEmpty(items)){
                        index.addAndGet(items.size());
                    }
                    pullItemResponse.setTaskName(taskName);
                    pullItemResponse.setItems(items);
                    pullItemResponse.setIndex(index.get());

                    responseMessage.setSuccess(true);
                    responseMessage.setStatus(ResponseStatus.SUCCESS.getCode());
                    responseMessage.setBodyObject(pullItemResponse);
                } catch (Exception e) {
                    log.error(e.getMessage(),e);
                    responseMessage.setSuccess(false);
                    responseMessage.setStatus(ResponseStatus.ERROR.getCode());
                    responseMessage.setMessage(e.getMessage());
                }
            }
            else {
                responseMessage.setSuccess(false);
                responseMessage.setStatus(ResponseStatus.NOT_MATCH_NODE.getCode());
                log.warn("response failed -> not match node.");
            }
        }
        else {
            responseMessage.setSuccess(false);
            responseMessage.setStatus(ResponseStatus.NO_HANDLER.getCode());
            log.warn("response failed -> no handler.");
        }
        return responseMessage;
    }
}
